package com.stock.sharding;

import com.google.common.collect.Range;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;
import org.joda.time.DateTime;

import java.util.Collection;
import java.util.Date;
import java.util.Optional;
import java.util.stream.Collectors;

//TODO 定义数据表的分片算法类：包含精准匹配数据表和范围匹配数据表
//      因为表是根据日期分表的，一月一个表，所以片键的类型是Date
public class CommonAlg4Tb implements PreciseShardingAlgorithm<Date>, RangeShardingAlgorithm<Date> {
    /**
     * TODO 精准匹配数据表的方法 cur_time 条件必须是 = 或者 in
     *
     * @param tbNames       所有可匹配数据源的集合 es:{202401..202412}
     * @param shardingValue 集成了SQL语句中的内容
     * @return
     */
    @Override
    public String doSharding(Collection<String> tbNames, PreciseShardingValue<Date> shardingValue) {
        //1.思路：根据传入的日期值，获取年份字符串
        //获取分片字段的名称colume
        String columnName = shardingValue.getColumnName();
        //获取逻辑表名称
        String logicTableName = shardingValue.getLogicTableName();
        //获取分片值
        Date value = shardingValue.getValue();
        //获取年份字符串
        String yearMonth = new DateTime(value).toString("yyyyMM");

        //2.获取匹配的数据源名称
        Optional<String> realTbName = tbNames.stream().filter(tbName -> tbName.endsWith(yearMonth)).findFirst();

        //3.判断是否有符合指定年份的数据源
        if (realTbName.isPresent()) {
            return realTbName.get();
        }
        return null;
    }

    /**
     * TODO 范围查询匹配数据源的方法 关键字：between and
     *
     * @param tbNames       所有可匹配数据源的集合 es:{202401..202412}
     * @param shardingValue 集成了SQL语句中的内容
     * @return
     */
    @Override
    public Collection<String> doSharding(Collection<String> tbNames, RangeShardingValue<Date> shardingValue) {
        //获取分片字段的名称colume
        String columnName = shardingValue.getColumnName();
        //获取逻辑表名称
        String logicTableName = shardingValue.getLogicTableName();

        //1.获取范围封装对象
        Range<Date> valueRange = shardingValue.getValueRange();

        //2.判断下限值
        if (valueRange.hasLowerBound()) {
            //获取下限日期
            Date lowerDate = valueRange.lowerEndpoint();
            //获取年月份
            int startYearMonth = Integer.parseInt(new DateTime(lowerDate).toString("yyyyMM"));

            tbNames = tbNames.stream().filter(tbName -> Integer.valueOf(tbName.substring(tbName.lastIndexOf("_") + 1)) >= startYearMonth)
                    .collect(Collectors.toList());
        }

        //3.判断上限值
        if (valueRange.hasUpperBound()) {
            //获取上限日期
            Date upperDate = valueRange.upperEndpoint();
            //获取年月份
            int endYearMonth = Integer.parseInt(new DateTime(upperDate).toString("yyyyMM"));

            tbNames = tbNames.stream().filter(tbName -> Integer.valueOf(tbName.substring(tbName.lastIndexOf("_") + 1)) <= endYearMonth)
                    .collect(Collectors.toList());
        }
        return tbNames;
    }
}
